import logging
from datetime import datetime
from random import random, randint

from a_mem import AgenticMemorySystem
from .dataset import LOCOMO
from .llm import generate_keywords_llm, answer_llm, QA_Evaluation
from .prompt import (
    generate_query_keywords,
    answer_adversarial_question_prompt,
    answer_multi_hop_question_prompt,
    answer_temporal_reasoning_question_prompt,
    answer_other_question_prompt
)


logger = logging.getLogger(__name__)


def session_datetime_to_timestamp(session_datetime: str) -> str:
    processed_datetime = datetime.strptime(session_datetime, "%I:%M %p on %d %B, %Y")
    output_datetime = datetime.strftime(processed_datetime, "%Y%m%d%H%M%S")
    return output_datetime


def timestamp_to_session_datetime(timestamp: str) -> str:
    processed_datetime = datetime.strptime(timestamp, "%Y%m%d%H%M%S")
    output_datetime = datetime.strftime(processed_datetime, "%I:%M %p on %d %B, %Y")
    return output_datetime


class LOCOMO_QA:
    def __init__(
            self,
            dataset: LOCOMO,
            memory_system: AgenticMemorySystem,
            only_one: bool = False
    ):
        self.dataset = dataset
        self.memory_system = memory_system
        self.memory_system.clear()
        self.random_integral = None
        if only_one:
            self.random_integral = randint(0, 9)

    def add_conversations(self) -> None:
        if self.random_integral:
            conversations = [self.dataset.conversations[self.random_integral]]
        else:
            conversations = self.dataset.conversations
        for conversation in conversations:
            for session in conversation:
                session_datetime = session.datetime
                output_datetime = session_datetime_to_timestamp(session_datetime)
                for turn in session.turns:
                    speaker, text = turn.speaker, turn.text
                    if blip_caption := turn.blip_caption:
                        content = f"Speaker {speaker} shares {blip_caption}, and says {text}"
                    else:
                        content = f"Speaker {speaker} says {text}"
                    self.memory_system.add_note(content, output_datetime)

    def answer_question(self, k: int = 5) -> list[QA_Evaluation]:
        answer_list = []
        if self.random_integral:
            qas = [self.dataset.qas[self.random_integral]]
        else:
            qas = self.dataset.qas
        for qa_list in qas:
            for qa in qa_list:
                question, evidence_list, category, final_answer = (qa.question, qa.evidence,
                                                                   qa.category, qa.final_answer)
                generate_keywords_prompt = generate_query_keywords.format(question=question)
                keywords = generate_keywords_llm.invoke(generate_keywords_prompt)
                documents, timestamps = self.memory_system.search(f"{keywords.keywords}", k)
                session_datetime = [ timestamp_to_session_datetime(timestamp)
                                     for timestamp in timestamps ]
                context = "\n".join([ f"talk start datetime: {session_datetime} {document}"
                                      for document, session_datetime
                                      in zip(documents, session_datetime) ])
                if category == 5:
                    if final_answer == "Yes":
                        other_answer = "No"
                    elif final_answer == "No":
                        other_answer = "Yes"
                    else:
                        other_answer = "Not mentioned in the conversation"
                    if random() < 0.5:
                        prompt = answer_adversarial_question_prompt.format(
                            context=context,
                            question=question,
                            answer1=other_answer,
                            answer2=final_answer
                        )
                    else:
                        prompt = answer_adversarial_question_prompt.format(
                            context=context,
                            question=question,
                            answer1=final_answer,
                            answer2=other_answer
                        )
                elif category == 2:
                    prompt = answer_multi_hop_question_prompt.format(
                        context=context,
                        question=question
                    )
                elif category == 3:
                    prompt = answer_temporal_reasoning_question_prompt.format(
                        context=context,
                        question=question
                    )
                else:
                    prompt = answer_other_question_prompt.format(
                        context=context,
                        question=question
                    )
                answer = answer_llm.invoke(prompt)
                logger.info(f"question: {question}, llm answer: {answer.answer}, "
                            f"real answer: {final_answer}")
                answer_list.append(QA_Evaluation(
                    llm_answer=answer.answer,
                    real_answer=final_answer,
                    category=category
                ))
        return answer_list
